#![feature(vec_into_raw_parts)]

pub mod c_api;
pub mod classes;
pub mod enums;
pub mod get_process_mitigation_command;
pub mod ntdll;
pub mod registry_manager;
pub mod set_process_mitigation_command;

pub use classes::*;
pub use enums::*;
pub use get_process_mitigation_command::*;
pub use ntdll::*;
pub use registry_manager::RegistryManager;
pub use set_process_mitigation_command::SetProcessMitigationsCommand;

use std::mem;
use std::path::Path;
use windows::Win32::Foundation::HANDLE;
use windows::Win32::System::Diagnostics::ToolHelp::{
    CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW, TH32CS_SNAPPROCESS,
};
use windows::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_INFORMATION, PROCESS_VM_READ};

#[derive(Debug, Clone)]
pub struct ProcessInfo {
    pub process_name: String,
    pub id: u32,
    pub handle: Option<HANDLE>,
}

impl ProcessInfo {
    pub fn get_handle(&mut self) -> Result<HANDLE, windows::core::Error> {
        if let Some(handle) = self.handle {
            Ok(handle)
        } else {
            unsafe {
                let handle =
                    OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, false, self.id)?;
                self.handle = Some(handle);
                Ok(handle)
            }
        }
    }
}

/// Represents the result of an operation that can either succeed or fail.
#[derive(Debug, Clone)]
pub struct AppResult<T> {
    pub is_success: bool,
    pub value: Option<T>,
    pub error: String,
}

impl<T> AppResult<T> {
    pub fn success(value: T) -> Self {
        Self {
            is_success: true,
            value: Some(value),
            error: String::new(),
        }
    }

    pub fn failure(error: &str) -> Self {
        Self {
            is_success: false,
            value: None,
            error: error.to_string(),
        }
    }

    pub fn is_failure(&self) -> bool {
        !self.is_success
    }

    pub fn match_result<R>(
        &self,
        on_success: impl FnOnce(&T) -> R,
        on_failure: impl FnOnce(&str) -> R,
    ) -> R {
        if self.is_success {
            if let Some(ref value) = self.value {
                on_success(value)
            } else {
                on_failure("Success result has no value")
            }
        } else {
            on_failure(&self.error)
        }
    }

    pub fn get_value_or_default(&self, default_value: T) -> T
    where
        T: Clone,
    {
        if self.is_success {
            self.value.clone().unwrap_or(default_value)
        } else {
            default_value
        }
    }
}

/// Result type for operations that don't return a value
#[derive(Debug, Clone)]
pub struct AppResultVoid {
    pub is_success: bool,
    pub error: String,
}

impl AppResultVoid {
    pub fn success() -> Self {
        Self {
            is_success: true,
            error: String::new(),
        }
    }

    pub fn failure(error: &str) -> Self {
        Self {
            is_success: false,
            error: error.to_string(),
        }
    }

    pub fn is_failure(&self) -> bool {
        !self.is_success
    }

    pub fn match_result<R>(
        &self,
        on_success: impl FnOnce() -> R,
        on_failure: impl FnOnce(&str) -> R,
    ) -> R {
        if self.is_success {
            on_success()
        } else {
            on_failure(&self.error)
        }
    }
}

fn parse_mitigation_numbers(numbers_str: &str) -> Result<Vec<MitigationOptions>, String> {
    if numbers_str.trim().is_empty() {
        return Ok(Vec::new());
    }

    let mut options = Vec::new();
    for num_str in numbers_str.split(',') {
        let num_str = num_str.trim();
        if !num_str.is_empty() {
            match num_str.parse::<u32>() {
                Ok(num) => match MitigationOptions::from_number(num) {
                    Ok(option) => options.push(option),
                    Err(e) => return Err(e),
                },
                Err(_) => return Err(format!("Invalid number format: {}", num_str)),
            }
        }
    }
    Ok(options)
}

fn parse_string_list(strings_str: &str) -> Vec<String> {
    if strings_str.trim().is_empty() {
        return Vec::new();
    }

    strings_str
        .split(',')
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect()
}

pub struct Program;

impl Program {
    /// Gets the current settings on all running instances of a specific process.
    pub fn get_from_running_process_name(
        running_process_name: &str,
    ) -> AppResult<Vec<ProcessMitigations>> {
        let process_name = Path::new(running_process_name)
            .file_stem()
            .and_then(|s| s.to_str())
            .unwrap_or(running_process_name);

        let processes = Self::get_processes_by_name(process_name);
        let mut output = Vec::new();

        for process in processes {
            match GetProcessMitigationCommand::get_policy_from_running_process(&process) {
                Ok(mitigation) => output.push(mitigation),
                Err(e) => {
                    return AppResult::failure(&format!(
                        "Failed to get policies from running process: {}",
                        e
                    ));
                }
            }
        }

        AppResult::success(output)
    }

    /// Gets the current system process mitigation defaults stored in the registry.
    pub fn get_system_policy() -> AppResult<AppMitigations> {
        match std::panic::catch_unwind(|| {
            let registry_manager = RegistryManager::new();
            registry_manager.get_system_policy()
        }) {
            Ok(Ok(policy)) => AppResult::success(policy),
            Ok(Err(error)) => AppResult::failure(&error),
            Err(_) => AppResult::failure("Failed to get system policy"),
        }
    }

    /// Gets all policies for all processes set in the registry.
    pub fn get_full_policies_for_all_processes() -> AppResult<Vec<AppMitigations>> {
        match std::panic::catch_unwind(|| {
            let registry_manager = RegistryManager::new();
            registry_manager.get_policy_list_from_registry(None)
        }) {
            Ok(Ok(policies)) => AppResult::success(policies),
            Ok(Err(error)) => AppResult::failure(&error),
            Err(_) => AppResult::failure("Failed to get full policies for all processes"),
        }
    }

    /// Gets the current settings in the registry for a specific process
    pub fn get_policy_from_registry_by_name(process_name: &str) -> AppResult<Vec<AppMitigations>> {
        match std::panic::catch_unwind(|| {
            let registry_manager = RegistryManager::new();
            registry_manager.get_policy_from_registry_by_name(process_name)
        }) {
            Ok(Ok(policies)) => AppResult::success(policies),
            Ok(Err(error)) => AppResult::failure(&error),
            Err(_) => AppResult::failure("Failed to get policy from registry by name"),
        }
    }

    /// Add mitigations for the Process
    pub fn add_mitigations_for_process(
        process_name: &str,
        disable_list: Option<&[MitigationOptions]>,
        enable_list: Option<&[MitigationOptions]>,
        eaf_modules_list: Option<&[String]>,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let mut registry_manager = RegistryManager::new();

        if enable_list.is_none() && disable_list.is_none() {
            return AppResultVoid::failure("Nothing was specified to enable or disable.");
        }

        let mut process_name = process_name.to_string();
        if process_name
            == Path::new(&process_name)
                .file_stem()
                .and_then(|s| s.to_str())
                .unwrap_or(&process_name)
        {
            process_name = format!("{}.exe", process_name).to_lowercase();
        }

        let from_registry_by_name =
            match registry_manager.get_policy_from_registry_by_name(&process_name) {
                Ok(policies) => policies,
                Err(error) => return AppResultVoid::failure(&error),
            };

        let mut app_mitigations = if from_registry_by_name.is_empty() {
            AppMitigations::new(&process_name)
        } else if from_registry_by_name.len() > 1 {
            return AppResultVoid::failure(
                "More than 1 mitigation policy found for the process. Use full path instead.",
            );
        } else {
            from_registry_by_name.into_iter().next().unwrap()
        };

        match SetProcessMitigationsCommand::set_enable_and_disable(
            &mut app_mitigations,
            &mut registry_manager,
            disable_list,
            enable_list,
            eaf_modules_list,
            is_force,
            is_remove,
            is_reset,
            false, // isSystemMode
        ) {
            Ok(_) => AppResultVoid::success(),
            Err(e) => {
                AppResultVoid::failure(&format!("Could not add mitigations for the process: {}", e))
            }
        }
    }

    /// Add mitigations to the System
    pub fn add_mitigations_to_system(
        disable_list: Option<&[MitigationOptions]>,
        enable_list: Option<&[MitigationOptions]>,
        eaf_modules_list: Option<&[String]>,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let mut registry_manager = RegistryManager::new();

        let mut app_mitigations = match registry_manager.get_system_policy() {
            Ok(policy) => policy,
            Err(_) => {
                let mut policy = AppMitigations::new("System");
                policy.source = "System Defaults".to_string();
                policy
            }
        };

        match SetProcessMitigationsCommand::set_enable_and_disable(
            &mut app_mitigations,
            &mut registry_manager,
            disable_list,
            enable_list,
            eaf_modules_list,
            is_force,
            is_remove,
            is_reset,
            true, // isSystemMode
        ) {
            Ok(_) => AppResultVoid::success(),
            Err(e) => {
                AppResultVoid::failure(&format!("Failed to add mitigations to system: {}", e))
            }
        }
    }

    /// Import mitigations from XML file
    pub fn import_xml_mitigations(file_path: &str) -> AppResultVoid {
        let result = SetProcessMitigationsCommand::import_mitigation(file_path);
        if result >= 0 {
            AppResultVoid::success()
        } else {
            AppResultVoid::failure(&format!("Error importing: {}", result))
        }
    }

    /// Validate Mitigations XML file
    pub fn validate_xml_mitigations(file_path: &str) -> AppResult<bool> {
        let mut result = false;
        let num = SetProcessMitigationsCommand::validate_xml_from_managed(file_path, &mut result);
        if num < 0 {
            AppResult::failure(&format!("Error validating: {}", num))
        } else {
            println!("Is Valid XML: {}", result);
            AppResult::success(result)
        }
    }

    fn get_processes_by_name(process_name: &str) -> Vec<ProcessInfo> {
        let mut processes = Vec::new();

        unsafe {
            let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
            if snapshot.is_err() {
                return processes;
            }
            let snapshot = snapshot.unwrap();

            let mut process_entry = PROCESSENTRY32W {
                dwSize: mem::size_of::<PROCESSENTRY32W>() as u32,
                ..Default::default()
            };

            if Process32FirstW(snapshot, &mut process_entry).is_ok() {
                loop {
                    let exe_file = String::from_utf16_lossy(&process_entry.szExeFile)
                        .trim_end_matches('\0')
                        .to_string();

                    if exe_file.eq_ignore_ascii_case(&format!("{}.exe", process_name))
                        || exe_file.eq_ignore_ascii_case(process_name)
                    {
                        processes.push(ProcessInfo {
                            process_name: exe_file,
                            id: process_entry.th32ProcessID,
                            handle: None,
                        });
                    }

                    if Process32NextW(snapshot, &mut process_entry).is_err() {
                        break;
                    }
                }
            }
        }

        processes
    }

    /// Add mitigations for the Process using comma-separated numbers
    pub fn add_mitigations_for_process_ex(
        process_name: &str,
        disable_list_numbers: &str,
        enable_list_numbers: &str,
        eaf_modules_list: &str,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let disable_options = match parse_mitigation_numbers(disable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing disable list: {}", e)),
        };

        let enable_options = match parse_mitigation_numbers(enable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing enable list: {}", e)),
        };

        let eaf_modules = parse_string_list(eaf_modules_list);
        let eaf_modules_ref = if eaf_modules.is_empty() {
            None
        } else {
            Some(eaf_modules.as_slice())
        };

        Self::add_mitigations_for_process(
            process_name,
            disable_options.as_deref(),
            enable_options.as_deref(),
            eaf_modules_ref,
            is_force,
            is_remove,
            is_reset,
        )
    }

    /// Add mitigations to the System using comma-separated numbers
    pub fn add_mitigations_to_system_ex(
        disable_list_numbers: &str,
        enable_list_numbers: &str,
        eaf_modules_list: &str,
        is_force: Option<&str>,
        is_remove: bool,
        is_reset: bool,
    ) -> AppResultVoid {
        let disable_options = match parse_mitigation_numbers(disable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing disable list: {}", e)),
        };

        let enable_options = match parse_mitigation_numbers(enable_list_numbers) {
            Ok(options) => {
                if options.is_empty() {
                    None
                } else {
                    Some(options)
                }
            }
            Err(e) => return AppResultVoid::failure(&format!("Error parsing enable list: {}", e)),
        };

        let eaf_modules = parse_string_list(eaf_modules_list);
        let eaf_modules_ref = if eaf_modules.is_empty() {
            None
        } else {
            Some(eaf_modules.as_slice())
        };

        Self::add_mitigations_to_system(
            disable_options.as_deref(),
            enable_options.as_deref(),
            eaf_modules_ref,
            is_force,
            is_remove,
            is_reset,
        )
    }

    /// Gets the current settings in the registry for a specific process and returns JSON
    pub fn get_policy_from_registry_by_name_ex(process_name: &str) -> AppResult<String> {
        match Self::get_policy_from_registry_by_name(process_name) {
            app_result if app_result.is_success => {
                if let Some(policies) = app_result.value {
                    let json_policies: Vec<AppMitigationsJson> =
                        policies.iter().map(|p| p.to_json()).collect();
                    match serde_json::to_string_pretty(&json_policies) {
                        Ok(json) => AppResult::success(json),
                        Err(e) => {
                            AppResult::failure(&format!("Failed to serialize to JSON: {}", e))
                        }
                    }
                } else {
                    AppResult::success("[]".to_string())
                }
            }
            app_result => AppResult::failure(&app_result.error),
        }
    }

    /// Gets the current system process mitigation defaults stored in the registry and returns JSON
    pub fn get_system_policy_ex() -> AppResult<String> {
        match Self::get_system_policy() {
            app_result if app_result.is_success => {
                if let Some(policy) = app_result.value {
                    let json_policy = policy.to_json();
                    match serde_json::to_string_pretty(&json_policy) {
                        Ok(json) => AppResult::success(json),
                        Err(e) => {
                            AppResult::failure(&format!("Failed to serialize to JSON: {}", e))
                        }
                    }
                } else {
                    AppResult::failure("No system policy found")
                }
            }
            app_result => AppResult::failure(&app_result.error),
        }
    }
}
